#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include "sm2_utils.h"
#include "sm3_utils.h"
#include "sm4_utils.h"

// 数字信封结构
typedef struct {
    uint8_t encrypted_key[256];  // SM2加密后的SM4密钥
    size_t encrypted_key_len;
    uint8_t ciphertext[1024];    // SM4加密后的数据
    size_t ciphertext_len;
    uint8_t hash[SM3_DIGEST_LENGTH]; // 原始数据的哈希值
} DigitalEnvelope;

// 发送方：创建数字信封
int create_digital_envelope(const uint8_t *data, size_t data_len,
                           const SM2_PublicKey *receiver_pub_key,
                           DigitalEnvelope *envelope) {
    int ret = 0;
    uint8_t sym_key[SM4_KEY_SIZE];
    
    // 1. 生成随机SM4密钥
    if ((ret = sm4_generate_key(sym_key)) != 0) {
        printf("生成SM4密钥失败\n");
        return ret;
    }
    
    // 2. 使用SM4密钥加密数据
    if ((ret = sm4_encrypt(data, data_len, sym_key, 
                          envelope->ciphertext, &envelope->ciphertext_len)) != 0) {
        printf("SM4加密数据失败\n");
        return ret;
    }
    
    // 3. 计算原始数据的SM3哈希值
    if ((ret = sm3_hash(data, data_len, envelope->hash)) != 0) {
        printf("计算SM3哈希值失败\n");
        return ret;
    }
    
    // 4. 使用接收方的SM2公钥加密SM4密钥
    if ((ret = sm2_encrypt(sym_key, SM4_KEY_SIZE, receiver_pub_key,
                          envelope->encrypted_key, &envelope->encrypted_key_len)) != 0) {
        printf("SM2加密密钥失败\n");
        return ret;
    }
    
    return 0;
}

// 接收方：打开数字信封
int open_digital_envelope(const DigitalEnvelope *envelope,
                         const SM2_PrivateKey *receiver_priv_key,
                         uint8_t *decrypted_data, size_t *decrypted_len) {
    int ret = 0;
    uint8_t sym_key[SM4_KEY_SIZE];
    size_t sym_key_len = SM4_KEY_SIZE;
    
    // 1. 使用SM2私钥解密SM4密钥
    if ((ret = sm2_decrypt(envelope->encrypted_key, envelope->encrypted_key_len,
                          receiver_priv_key, sym_key, &sym_key_len)) != 0) {
        printf("SM2解密密钥失败\n");
        return ret;
    }
    
    // 2. 使用SM4密钥解密数据
    if ((ret = sm4_decrypt(envelope->ciphertext, envelope->ciphertext_len,
                          sym_key, decrypted_data, decrypted_len)) != 0) {
        printf("SM4解密数据失败\n");
        return ret;
    }
    
    // 3. 验证哈希值
    if ((ret = sm3_verify(decrypted_data, *decrypted_len, envelope->hash)) != 0) {
        printf("哈希验证失败，数据可能被篡改\n");
        return ret;
    }
    
    return 0;
}

int main() {
    int ret = 0;
    
    // 测试数据
    const char *test_data = "这是一条测试消息，用于演示数字信封协议。";
    size_t test_data_len = strlen(test_data);
    
    // 接收方生成SM2密钥对
    SM2_KeyPair receiver_keypair;
    if ((ret = sm2_generate_keypair(&receiver_keypair)) != 0) {
        printf("生成SM2密钥对失败\n");
        return ret;
    }
    
    printf("接收方SM2密钥对生成成功\n");
    
    // 创建数字信封
    DigitalEnvelope envelope;
    if ((ret = create_digital_envelope((const uint8_t *)test_data, test_data_len,
                                      &receiver_keypair.public_key, &envelope)) != 0) {
        printf("创建数字信封失败\n");
        return ret;
    }
    
    printf("发送方创建数字信封成功\n");
    printf("- 加密后的SM4密钥长度: %zu 字节\n", envelope.encrypted_key_len);
    printf("- 加密后的数据长度: %zu 字节\n", envelope.ciphertext_len);
    
    // 打开数字信封
    uint8_t decrypted_data[1024];
    size_t decrypted_len = 0;
    
    if ((ret = open_digital_envelope(&envelope, &receiver_keypair.private_key,
                                    decrypted_data, &decrypted_len)) != 0) {
        printf("打开数字信封失败\n");
        return ret;
    }
    
    printf("接收方打开数字信封成功\n");
    printf("- 解密后的数据长度: %zu 字节\n", decrypted_len);
    printf("- 解密后的数据内容: %.*s\n", (int)decrypted_len, decrypted_data);
    
    // 验证解密后的数据是否与原始数据一致
    if (decrypted_len == test_data_len && 
        memcmp(decrypted_data, test_data, test_data_len) == 0) {
        printf("数据验证成功，解密后的数据与原始数据一致\n");
    } else {
        printf("数据验证失败，解密后的数据与原始数据不一致\n");
    }
    
    return 0;
}